package cn.liziguo.randomfilegenerator;

import java.io.File;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.util.concurrent.ThreadLocalRandom;
import java.util.function.DoubleConsumer;

/**
 * @author Liziguo
 * @date 2024-11-04
 */
public class RandomFileFactory {
    public static final int BUFF_SIZE = 1024 * 1024;

    /**
     * 总大小
     */
    protected final long total;
    /**
     * 文件生成目录
     */
    protected final File dir;
    /**
     * 单个文件大小
     */
    protected final long fileSize;
    /**
     * 文件数量
     */
    protected final int count;
    /**
     * 当进度条发送变化回调
     */
    protected final DoubleConsumer onProgress;
    /**
     * 每秒速度计算回调
     */
    protected final DoubleConsumer onSpeed;
    /**
     * 程序结束回调
     */
    protected final Runnable onFinish;
    /**
     * 已创建的大小
     */
    protected long size = 0;
    /**
     * 是否正在运行
     */
    protected boolean running;
    /**
     * 用于计算速度 开始时的时间
     */
    protected long start = 0;
    /**
     * 用于计算速度 开始时的指针位置
     */
    protected long startPos = 0;

    public RandomFileFactory(File dir, long fileSize, int count, DoubleConsumer onProgress, DoubleConsumer onSpeed, Runnable onFinish) {
        total = fileSize * count;
        this.dir = dir;
        this.fileSize = fileSize;
        this.count = count;
        this.onProgress = onProgress;
        this.onSpeed = onSpeed;
        this.onFinish = onFinish;
    }

    protected String getFileName(int index) {
        return System.currentTimeMillis() + "_" + index;
    }

    public void create() throws IOException {
        running = true;
        try {
            File parentFile = dir;
            if (!parentFile.exists()) {
                parentFile.mkdirs();
            }
            start = System.nanoTime();
            if (fileSize > 0) {
                for (int i = 0; i < count && running; i++) {
                    createFile(new File(parentFile, getFileName(i)), fileSize);
                }
            } else {
                for (int i = 0; i < count && running; i++) {
                    createFile(new File(parentFile, getFileName(i)), fileSize);
                    onProgress.accept((double) i / count);
                }
            }
        } finally {
            running = false;
            onFinish.run();
        }
    }

    public void stop() {
        running = false;
    }

    protected void createFile(File file, long fileSize) throws IOException {
        ThreadLocalRandom random = ThreadLocalRandom.current();
        try (RandomAccessFile rw = new RandomAccessFile(file, "rw")) {
            byte[] bytes = new byte[BUFF_SIZE];
            long pos = 0;
            while (pos < fileSize && running) {
                random.nextBytes(bytes);
                int len;
                if (fileSize - pos > Integer.MAX_VALUE) {
                    len = BUFF_SIZE;
                    rw.write(bytes, 0, BUFF_SIZE);
                } else {
                    len = Math.min(BUFF_SIZE, (int) (fileSize - pos));
                    rw.write(bytes, 0, len);
                }
                pos += len;
                size += len;
                long end = System.nanoTime();
                long duration = end - start;
                if (duration > 1_000_000_000) {
                    double d = duration / 1_000_000_000.0;
                    double speed = (size - startPos) / d;
                    onSpeed.accept(speed);
                    start = System.nanoTime();
                    startPos = size;
                }
                onProgress.accept((double) size / total);
            }
        }
    }
}
